package com.study.spring_ai.tool;

import java.time.LocalDateTime;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.context.i18n.LocaleContextHolder;

/**
 * 声明式工具
 */
public class DateTimeTools {

    @Tool(description = "Get the current date and time in the user's timezone")
    String getCurrentDateTime() {
        return LocalDateTime.now().atZone(LocaleContextHolder.getTimeZone().toZoneId()).toString();
    }

    @Tool(description = "Get the current date and time in the user's timezone")
    String printCurrentDateTime(@ToolParam(description = "The timezone to use") String currentTimeZone) {
        return LocalDateTime.now().atZone(LocaleContextHolder.getTimeZone().toZoneId()).toString();
    }


    public static void main(String[] args) {
        DateTimeTools dateTimeTools = new DateTimeTools();
        ChatClient.builder(OllamaChatModel.builder().ollamaApi(OllamaApi.builder().build()).build())
                .defaultTools(dateTimeTools);
    }

    /**
     * 手动控制工具执行
     */
    public void test() {
        //1、构建ToolCallback
        DateTimeTools dateTimeTools = new DateTimeTools();
        OllamaChatModel chatModel =  OllamaChatModel.builder().ollamaApi(OllamaApi.builder().build()).build();
        ToolCallingManager toolCallingManager = ToolCallingManager.builder().build();
        ToolCallback[] from = ToolCallbacks.from(dateTimeTools);
        //2、构建ChatOptions 组装提示词
        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(from)
                .internalToolExecutionEnabled(false)
                .build();
        Prompt prompt = new Prompt("Tell me more about the customer with ID 42", chatOptions);
        //3、调用模型
        ChatResponse chatResponse = chatModel.call(prompt);
        //4、手动控制工具调用
        while (chatResponse.hasToolCalls()) {

            ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, chatResponse);

            prompt = new Prompt(toolExecutionResult.conversationHistory(), chatOptions);

            chatResponse = chatModel.call(prompt);
        }

        System.out.println(chatResponse.getResult().getOutput().getText());

    }


}